/*
    Copyright 2006-2012 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Contains classes for grain species.

#ifndef __grain__
#define __grain__

#include <iostream>
#include "blitz/array.h"
#include "blackbody.h"
#include "solvert.h"
#include "mcrx-types.h"
#include "mcrx-units.h"

#include "mono_poly_abstract.h"
#include <vector>
#include "misc.h"
#include <string>
#include "threadlocal.h"
#include "interpolatort.h"
#include "vecops.h"
#include "constants.h"
#include "mcrx-debug.h"

namespace mcrx {
  class grain_data;
  class emitting_grain;
  class template_emission_grain;
  class Brent_PAH_grain;
  class thermal_equilibrium_grain;
};

int main(int, char**); // for testing purposes


class mcrx::grain_data {
private:
  /// Grain absorption cross-section as a function of (size,
  /// wavelength). Used for absorption calculations.
  array_2 sigma_; 
  /// Grain absorption cross-section as a function of (size,
  /// wavelength) on the wavelength grid for emission. Used for
  /// emission calculations.
  array_2 esigma_; 
  array_2 sigma_sca_; ///< Grain scattering cross section.
  array_2 g_; ///< Grain \<cos theta\> of scattered light.

  /// Size vector.
  std::vector<T_float> sizes_;
  array_1 asizes_; ///< Array view of size vector.
  array_1 delta_size_;

  /// Wavelength vector.
  std::vector<T_float> lambda_;
  array_1 alambda_; ///< Array view of wavelength vector.
  array_1 elambda_; ///< The wavelength vector for emission calculations.
  array_1 invelambda_; ///< The inverse wavelength vector for emission calculations.
  array_1 delta_lambda_;

  /// Units of cross-sections.
  T_unit_map units_;

  /** Loads opacities from file. By default all size bins are loaded,
      but min and max sizes can be specified if so desired. */
  void load_file (const std::string&, T_float, T_float);

protected:
  /** Set up internal variables. This is virtual so it can be
      overridden in derived classes, that must then make sure to call
      the base class setup as part of it. This function is called
      after construction and wavelength resampling. */
  virtual void setup();

  /** Returns the luminosity absorbed by species s in the specified
      radiation intensity. This is calculated by integrating the
      intensity times the absorption cross section. */
  template<typename T>
  T_float calculate_heating (int s, const blitz::ETBase<T>& i) const {
    const T_float hh= 4.0*constants::pi*
      integrate_quantity(sigma()(s, blitz::Range::all())*i,
			 alambda_, false);
    return hh;
  };

  /// \todo clean up
  /* 
  template<typename T>
  blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprBinaryOp<blitz::_bz_ArrayExprConstant<double>, typename T_integrate_quantities_lin<blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprBinaryOp<typename blitz::FastArrayIterator<double, 2>, typename blitz::asExpr<T>::T_expr, blitz::Multiply<double, double> > > >::T_result,blitz::Multiply<double, double> > >
  calculate_heating (const blitz::ETBase<T>& i) const {
    assert(0);
    return (4.0*constants::pi)*
      //RONG
      integrate_quantities_lin(sigma()*i, alambda_);
  };
  */

public:
  /** Constructor takes a file containing cross section information
      and loads that data. By default all size bins are loaded, but
      min and max sizes can be specified if so desired. */
  grain_data (const std::string& file_name, T_float minsize=0, 
	      T_float maxsize=blitz::huge(T_float()));
  grain_data (const grain_data& g) :
    // need to make sure all array members are real copies
    sigma_(make_thread_local_copy(g.sigma_)),
    esigma_(make_thread_local_copy(g.esigma_)),
    sigma_sca_(make_thread_local_copy(g.sigma_sca_)),
    g_(make_thread_local_copy(g.g_)),
    sizes_(g.sizes_),
    // asizes_ references sizes_
    asizes_(reference_vector(sizes_)),
    delta_size_(make_thread_local_copy(g.delta_size_)),
    lambda_(g.lambda_), 
    // alambda_ references lambda_
    alambda_(reference_vector(lambda_)),
    elambda_(make_thread_local_copy(g.elambda_)),
    invelambda_(make_thread_local_copy(1./elambda_)),
    units_(g.units_) {};
  virtual ~grain_data () {};

  /** Resamples the opacity data on to the specified wavelength vector.  */
  void resample (const array_1& lambda, const array_1& elambda=array_1());

  // cross-section retrieval functions
  array_1 sigma (int i) const {return sigma_(i, blitz::Range::all ());};
  array_1 esigma (int i) const {return esigma_(i, blitz::Range::all ());};
  const array_2& sigma () const {return sigma_;};
  const array_2& esigma () const {return esigma_;};
  const array_2& sigma_sca () const {return sigma_sca_;};
  const array_2& g () const {return g_;};

  // Size retrieval functions
  const std::vector<T_float>& sizes () const {return sizes_;};
  int n_size() const {return sizes_.size();};
  const array_1& asizes () const {return asizes_;};
  const array_1& delta_size () const {return delta_size_;};

  // Wavelength retrieval functions
  const std::vector<T_float>& lambda ()const {return lambda_;};
  int n_lambda() const { return lambda_.size();};
  int n_elambda() const { return elambda_.size();};
  const array_1& alambda ()const {return alambda_;};
  const array_1& elambda ()const {return elambda_;};
  const array_1& invelambda ()const {return invelambda_;};

  const T_unit_map& units() const {return units_;};
};



/** This is the abstract base class for the classes of grains for
    which the SED can be calculated. */
class mcrx::emitting_grain : public grain_data {
public:
  emitting_grain (const std::string& file_name, T_float minsize=0, 
	      T_float maxsize=blitz::huge(T_float())) : 
    grain_data(file_name, minsize, maxsize) {};
  // copy constructor implicitly generated

  /** This function calculates the emission spectrum of the grains
      based on the specified heating intensity and dust mass in the
      cells and the specified size distribution. mdust is the mass of
      dust in the cells, while dn is the size distribution expressed
      as number of dust grains per size bin per unit dust mass. The
      calculation is done in parallel or on a GPU, if applicable. For
      efficiency, the result is written to the existing array sed, or
      if add_to_sed is true, ADDED to existing data in the sed
      array. */
  template<typename T>
  void calculate_SED_from_intensity (const blitz::ETBase<T>& intensity,
				     const array_1& mdust,
				     const array_1& dn,
				     array_2 sed, 
				     bool add_to_sed,
				     array_2* temp) const;

  virtual boost::shared_ptr<emitting_grain> clone() const=0;
};


/** The template_emission_grain class implement grain which absorbs
    radiation with a specified absorption opacity, but then emits with
    a fixed template spectrum.  */  
class mcrx::template_emission_grain : public emitting_grain {
protected:
  /// The emission template, normalized for unit luminosity.
  array_1 template_emission_;

  // This class is not meant for end-users, so the constructors are not public.
  template_emission_grain (const std::string& file_name, T_float minsize=0, 
			   T_float maxsize=blitz::huge(T_float())) :
    emitting_grain(file_name, minsize, maxsize) {}
  template_emission_grain (const template_emission_grain& g) :
    emitting_grain(g),
    template_emission_(make_thread_local_copy(g.template_emission_))
  {};

public:

  /** This function calculates the emission spectrum of the grains
      based on the specified heating intensity and dust mass in the
      cells and the specified size distribution.  For efficiency, the
      result is written to the existing array sed, or if add_to_sed is
      true, added to existing data in the sed array. */
  template<typename T>
  void calculate_SED_from_intensity_virtual (const blitz::ETBase<T>& intensity,
					     const array_1& mdust,
					     const array_1& dn,
					     array_2 sed, 
					     bool add_to_sed,
					     array_2* temp) const;
 
};

/** This class implements Brent Groves's PAH template emission.  */
class mcrx::Brent_PAH_grain: public template_emission_grain {
protected:
  friend int ::main(int, char**); // for testing purposes

  static T_float x0_[], f0_[], sigma_[];

  /** Calculate the PAH template on the current wavelength grid.  */
  virtual void setup ();
  
public:
  Brent_PAH_grain (const std::string& opacity_file):
    template_emission_grain::template_emission_grain (opacity_file) {
    template_emission_grain::setup();};
  Brent_PAH_grain (const Brent_PAH_grain& g):
    template_emission_grain::template_emission_grain (g) {};
  virtual boost::shared_ptr<emitting_grain> 
  clone() const {
    return boost::shared_ptr<emitting_grain>
      (new Brent_PAH_grain(*this));};
};


/** Calculates the emission SED of the template grains.  This is much
   simpler than for the thermal grains, we just sum up the absorbed
   energy for each of the grains and multiply by the emission
   template. For efficiency's sake, we actually add up the cross
   sections and then integrate, instead of the other way. This
   expression is verified to do the same as calling calculate_heating
   for each species. */
template<typename T>
void mcrx::template_emission_grain::
calculate_SED_from_intensity_virtual (const blitz::ETBase<T>& intensity,
				      const array_1& mdust,
				      const array_1& dn,
				      array_2 sed, 
				      bool add_to_sed,
				      array_2* temp) const
{
  std::cout << "Calculating grain template emission" << std::endl;

  using namespace blitz;
  const size_t nc = mdust.size();

  if(temp)
    // temp does not make sense for the template grain.
    *temp = blitz::quiet_NaN(T_float());

  array_1 sigma_tot(sum(sigma()(tensor::j, tensor::i)*dn(tensor::j), tensor::j));
  if(!add_to_sed)
    sed=0;
  array_1 absorption
    (integrate_quantities_lin(sigma_tot(tensor::j)*intensity, alambda()));
  ASSERT_ALL(absorption>=0);
  sed += (4.0*constants::pi)*template_emission_(tensor::j) * 
    absorption(tensor::i) * mdust(tensor::i);
}
    


#endif
